from datetime import datetime, timedelta

from sqlalchemy.orm import Session

from app.bigdata.models.impala_query_info import ImpalaQueryInfo
from app.bigdata.utils import collection


def compute_stats_handler(db: Session, start_datatime: str, end_datatime: str):
    res = []
    items = ImpalaQueryInfo.get_statement_by_range_datetime(db, start_datatime, end_datatime)
    for item in items:
        date_obj = datetime.strptime(item.get("start_time"), '%Y-%m-%d %H:%M:%S')
        new_date_obj = (date_obj - timedelta(minutes=50)).strftime("%Y-%m-%d %H:%M:%S")
        table_name = item.get("statement").split(" ")[-1]
        data = ImpalaQueryInfo.get_statement_by_range_datetime(db, new_date_obj, item.get("start_time"),
                                                               statement=f'insert overwrite%{table_name}%')
        data.extend(ImpalaQueryInfo.get_statement_by_range_datetime(db, start_datatime, item.get("start_time"),
                                                                    statement=f'create%table%{table_name}%'))

        message = {"produced_table_query_id": "", "produced_table_statement": "", "rows_inserted": 0, "hdfs_bytes_written": 0}
        if data:
            statement = data[-1].get("statement").replace('\n', ' ')
            statement = ' '.join(statement.split())

            message = {"produced_table_query_id": data[-1].get("query_id"),
                       "produced_table_statement": statement[0:60],
                       "rows_inserted": data[-1]['query_response']['attributes'].get("rows_inserted"),
                       "hdfs_bytes_written": data[-1]['query_response']['attributes'].get("hdfs_bytes_written")}

        data = {"start_time": item.get("start_time"),
                "query_id": item.get("query_id"),
                "statement": ' '.join(item.get("statement").replace('\n', ' ').split()),
                "duration_millis": item.get("duration_millis"),
                'create_date': collection(item.get('start_time'), cst=False, dtype='cst_str').split(' ')[0]}
        data.update(message)
        res.append(data)
    return res


if __name__ == '__main__':
    from common.database import SessionLocal
    from pprint import pprint

    session = SessionLocal()
    pprint(compute_stats_handler(session, "2024-01-18 00:00:00", "2024-01-18 23:59:59"))
